#include "kcentergreedyint8.h"
#include "pccore.h"
#include "pcfeature.h"
#include "opencv2/imgproc.hpp"
#include "source/rknntool.h"
#include <QDebug>
#include <QElapsedTimer>
#include <rknn_api.h>
#include <QMutexLocker>

PCFeature::PCFeature(QObject *parent): QObject(parent), ctx(0), modelData(nullptr), modelDataSize(0), canRun(false) {}

PCFeature::~PCFeature() {
    unloadModel();
}

int PCFeature::loadModel(const QString &modelPath, QVariantMap &modelDetail)
{
    QMutexLocker locker(&mutex);
    qInfo() << "Loading model:" << modelPath;

    // Unload if already loaded
    if (ctx) unloadModel();

    FILE *fp = fopen(modelPath.toStdString().c_str(), "rb");
    if (!fp) {
        errorString = "无法打开模型文件: 文件不存在";
        return -1;
    }

    // Load model data
    fseek(fp, 0, SEEK_END);
    modelDataSize = ftell(fp);
    rewind(fp);
    modelData = (unsigned char*)malloc(modelDataSize);
    if (!modelData) {
        fclose(fp);
        errorString = "模型加载失败: 内存分配错误";
        return -1;
    }
    fread(modelData, 1, modelDataSize, fp);
    fclose(fp);

    // Initialize RKNN model
    int ret = rknn_init(&ctx, modelData, modelDataSize, 0, nullptr);
    if (ret < 0) {
        free(modelData);
        errorString = "模型初始化失败: 请检查模型文件的有效性";
        return -1;
    }

    // Query model details
    if (queryModelInfo(modelDetail) < 0) {
        free(modelData);
        return -1;
    }
    qInfo() << "PCFeature::loadModel";
    canRun = true;
    return 0;
}

int PCFeature::unloadModel()
{
    QMutexLocker locker(&mutex);
    if (!ctx) return 0;

    if (input_mems[0]) rknn_destroy_mem(ctx, input_mems[0]);
    if (output_mems[0]) rknn_destroy_mem(ctx, output_mems[0]);
    rknn_destroy(ctx);

    free(modelData);
    ctx = 0;
    modelData = nullptr;
    canRun = false;

    qInfo() << "Model successfully unloaded.";
    return 0;
}

int PCFeature::queryModelInfo(QVariantMap &modelDetail)
{
    // Query SDK version
    int ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(version));
    if (ret < 0) return -1;

    // Query input and output attributes
    memset(inputAttrs, 0, sizeof(inputAttrs));
    inputAttrs[0].index = 0;
    if (rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &inputAttrs[0], sizeof(inputAttrs[0])) < 0) return -1;
    dump_tensor_attr(&(inputAttrs[0]));

    memset(outputAttrs, 0, sizeof(outputAttrs));
    outputAttrs[0].index = 0;
    if (rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &outputAttrs[0], sizeof(outputAttrs[0])) < 0) return -1;
    dump_tensor_attr(&(outputAttrs[0]));
    // Set IO memory
    input_mems[0] = rknn_create_mem(ctx, inputAttrs[0].size);
    output_mems[0] = rknn_create_mem(ctx, outputAttrs[0].size);
    if (!input_mems[0] || !output_mems[0]) return -1;

    rknn_set_io_mem(ctx, input_mems[0], &inputAttrs[0]);
    rknn_set_io_mem(ctx, output_mems[0], &outputAttrs[0]);
    //NHWC
    input_size = cv::Size(inputAttrs[0].dims[2], inputAttrs[0].dims[1]);
    feture_size = cv::Size(outputAttrs[0].dims[0], outputAttrs[0].dims[1]);

    modelDetail["type"] = get_type_string(inputAttrs[0].type);
    modelDetail["size"] = static_cast<quint64>(modelDataSize);
    return 0;
}

void *PCFeature::detect(cv::Mat &baseImage)
{
    QMutexLocker locker(&mutex);
    if (!canRun) {
        qWarning() << "Model not loaded, cannot detect.";
        return nullptr ;
    }
    cv::Mat preprocessedImage = preprocessImage(baseImage);
    cv::Mat inputImage;
    QString modelType = modelDetal_["type"].toString();

    memcpy(input_mems[0]->virt_addr, preprocessedImage.data, inputAttrs[0].size);

    // QElapsedTimer timer;
    // timer.start();
    if (rknn_run(ctx, nullptr) < 0) return nullptr;
    //qInfo() << "PCFeature::detect:" << timer.nsecsElapsed() << "ns";
    //int8_t *feture = static_cast<int8_t*>(output_mems[0]->virt_addr);

    //qInfo() << "Total number of elements:"<< output_mems[0]->size;
    return output_mems[0]->virt_addr;
}

const char* PCFeature::get_type_string(int type) {
    switch (type) {
    case RKNN_TENSOR_FLOAT16: return "FP16";
    case RKNN_TENSOR_FLOAT32: return "FP32";
    case RKNN_TENSOR_INT8: return "INT8";
    default: return "Unknown";
    }
}

cv::Mat PCFeature::preprocessImage(const cv::Mat &inputImage)
{
    // 检查图像是否需要预处理
    if (inputImage.size() == modelInputSize && inputImage.channels() == 3) {
        return inputImage; // 如果图像已经符合要求，直接返回
    }

    // Step 1: Resize to 256x256
    cv::Mat resizedImage;
    cv::resize(inputImage, resizedImage, modelInputSize, 0, 0, cv::INTER_AREA);

    // // Step 2: Center crop to 224x224
    // int cropHeight = static_cast<int>(modelInputSize.height * (224.0 / 256));
    // int cropWidth = static_cast<int>(modelInputSize.width * (224.0 / 256));
    // int startX = (modelInputSize.width - cropWidth) / 2;
    // int startY = (modelInputSize.height - cropHeight) / 2;

    // cv::Rect cropRegion(startX, startY, cropWidth, cropHeight);
    // cv::Mat croppedImage = resizedImage(cropRegion);

    cv::Mat rgbImage;
    cv::cvtColor(resizedImage, rgbImage, cv::COLOR_BGR2RGB);

    return rgbImage;
}

